package com.lu.stack.gateway.filter.factory;

import cn.hutool.core.collection.CollectionUtil;
import cn.hutool.core.util.ObjectUtil;
import cn.hutool.core.util.StrUtil;
import cn.hutool.log.level.Level;
import com.lu.stack.common.model.response.R;
import com.lu.stack.gateway.utils.HttpResponseUtils;
import lombok.Setter;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Qualifier;
import org.springframework.cloud.gateway.filter.GatewayFilter;
import org.springframework.cloud.gateway.filter.factory.AbstractGatewayFilterFactory;
import org.springframework.cloud.gateway.filter.ratelimit.KeyResolver;
import org.springframework.cloud.gateway.filter.ratelimit.RateLimiter;
import org.springframework.cloud.gateway.route.Route;
import org.springframework.cloud.gateway.support.ServerWebExchangeUtils;
import org.springframework.http.HttpStatus;
import org.springframework.stereotype.Component;

import java.util.List;
import java.util.Map;

/**
 * @author zhanglu
 * @project lu-cloud-stack
 * @description 自定义请求限流
 * @date 2023/2/6 16:49
 */
@Slf4j
@Component
public class CustomRequestRateLimiterGatewayFilterFactory extends AbstractGatewayFilterFactory<CustomRequestRateLimiterGatewayFilterFactory.Config> {

    private final RateLimiter defaultRateLimiter;

    private final KeyResolver defaultKeyResolver;

    public CustomRequestRateLimiterGatewayFilterFactory(RateLimiter defaultRateLimiter,
                                                        @Qualifier("ipKeyResolver") KeyResolver defaultKeyResolver) {
        super(Config.class);
        this.defaultRateLimiter = defaultRateLimiter;
        this.defaultKeyResolver = defaultKeyResolver;
    }

    @Override
    public GatewayFilter apply(CustomRequestRateLimiterGatewayFilterFactory.Config config) {
        KeyResolver resolver = getOrDefault(config.keyResolver, defaultKeyResolver);
        RateLimiter<Object> limiter = defaultRateLimiter;
        return (exchange, chain) -> resolver.resolve(exchange)
                .flatMap(key -> {
                    Route route = exchange.getAttribute(ServerWebExchangeUtils.GATEWAY_ROUTE_ATTR);
                    String ip = exchange.getRequest().getRemoteAddress().getAddress().getHostAddress();
                    String uri = exchange.getRequest().getPath().toString();
                    if (ObjectUtil.isEmpty(key)) {
                        if (config.logLevels.stream().anyMatch(o -> o == Level.ERROR)) {
                            log.error(StrUtil.format("无效的访问令牌,应用:{},ip:{},url:{},key:{}", route.getId(), ip, uri, key));
                        }
                        return HttpResponseUtils.write(exchange.getResponse(), config.statusCode, R.failed().msg("访问已限流，无效的访问令牌"));
                    }
                    return limiter.isAllowed(route.getId(), key).flatMap(response -> {
                        for (Map.Entry<String, String> header : response.getHeaders().entrySet()) {
                            exchange.getResponse().getHeaders().add(header.getKey(), header.getValue());
                        }
                        if (response.isAllowed()) {
                            return chain.filter(exchange);
                        }
                        if (config.logLevels.stream().anyMatch(o -> o == Level.ERROR)) {
                            log.error(StrUtil.format("网关限流,应用:{},ip:{},url:{},key:{}", route.getId(), ip, uri, key));
                        }
                        return HttpResponseUtils.write(exchange.getResponse(), config.statusCode, R.failed().msg("访问已限流，请稍候重试"));
                    });
                });
    }

    private <T> T getOrDefault(T configValue, T defaultValue) {
        return (configValue != null) ? configValue : defaultValue;
    }

    @Setter
    public static class Config {
        /**
         * 限流key规则
         */
        private KeyResolver keyResolver;

        /**
         * 限流返回状态码
         */
        private HttpStatus statusCode = HttpStatus.TOO_MANY_REQUESTS;

        /**
         * 允许打印的日志级别
         */
        private List<Level> logLevels = CollectionUtil.newArrayList(Level.INFO, Level.ERROR);

    }

}
